import numpy as np
import matplotlib.pyplot as plt
import torch
from sklearn.decomposition import PCA
import colorsys
# 创建一个新的掩码，将-1转换为黑色
def visualize_mask(mask):
    # 创建一个RGB图像，初始化为全0（黑色）
    vis_mask = np.zeros((*mask.shape, 3))
    
    # 找到不是-1的位置
    valid_mask = mask != -1
    
    # 为非-1的位置分配随机颜色
    unique_labels = np.unique(mask[valid_mask])
    colors = plt.cm.get_cmap('tab20')(np.linspace(0, 1, len(unique_labels)))[:, :3]
    
    for i, label in enumerate(unique_labels):
        vis_mask[mask == label] = colors[i]
    
    return vis_mask
def id2rgb(id, max_num_obj=256):
        if not 0 <= id <= max_num_obj:
            raise ValueError("ID should be in range(0, max_num_obj)")
        # Convert the ID into a hue value
        golden_ratio = 1.6180339887
        h = ((id * golden_ratio) % 1)           # Ensure value is between 0 and 1
        s = 0.5 + (id % 2) * 0.5       # Alternate between 0.5 and 1.0
        l = 0.5
        # Use colorsys to convert HSL to RGB
        rgb = np.zeros((3, ), dtype=np.uint8)
        if id==0:   #invalid region
            return rgb
        r, g, b = colorsys.hls_to_rgb(h, l, s)
        rgb[0], rgb[1], rgb[2] = int(r*255), int(g*255), int(b*255)

        return rgb

def visualize_obj(objects):
        rgb_mask = np.zeros((*objects.shape[-2:], 3), dtype=np.uint8)
        all_obj_ids = np.unique(objects)
        for id in all_obj_ids:
            colored_mask = id2rgb(id)
            rgb_mask[objects == id] = colored_mask
        return rgb_mask

seg_map = np.load('/mnt/c/Users/cyt/Downloads/mask/mask_0.npy')
feature_map = np.load('/mnt/c/Users/cyt/Downloads/features_dim16/feature_0.npy')
'''
print(np.unique(seg_map))
print(feature_map.shape,seg_map.shape)
#vis_mask = visualize_mask(seg_map)
vis_mask = visualize_obj(seg_map)
# 创建图像显示
plt.figure(figsize=(10, 10))
plt.imshow(vis_mask)
plt.axis('off')  # 关闭坐标轴
plt.show()
'''
feature_map = torch.from_numpy(feature_map)
seg_map = torch.from_numpy(seg_map)
print(seg_map.shape)
# y,x [680,1200]
y, x = torch.meshgrid(torch.arange(0, 680), torch.arange(0, 1200))
# y,x [816000,1]
x = x.reshape(-1, 1)
y = y.reshape(-1, 1)
#结果: 得到一个[816000]的一维张量，包含了每个像素位置对应的分割标签
seg = seg_map[y, x].squeeze(-1).long()
mask = seg != -1
point_feature1 = feature_map[seg].squeeze(0)
mask = mask.reshape(1, 680, 1200)
point_feature = point_feature1.reshape(680, 1200, -1).permute(2, 0, 1)
print(point_feature.shape, mask.shape)
result = point_feature*mask

feature_numpy = result.permute(1, 2, 0).reshape(-1, result.shape[0]).numpy()
valid_features = feature_numpy[mask.reshape(-1) == 1]
# 执行PCA降维到3维
pca = PCA(n_components=3)
reduced_features = pca.fit_transform(valid_features)
reduced_features = (reduced_features - reduced_features.min(axis=0)) / (reduced_features.max(axis=0) - reduced_features.min(axis=0))
vis_img = np.zeros((680, 1200, 3))
valid_indices = np.where(mask.reshape(-1) == 1)[0]
vis_img.reshape(-1, 3)[valid_indices] = reduced_features
plt.figure(figsize=(15, 8))
plt.imshow(vis_img)
plt.axis('off')
plt.title('PCA Visualization of Features')
plt.show()
